package ratelimit

import (
	"sync"
	"time"
)

type TokenBucket struct {
	m sync.Mutex
	c int64
	t int64
}

func NewTokenBucket() *TokenBucket {
	return &TokenBucket{c: 1, t: time.Now().UnixNano()}
}

func (b *TokenBucket) Acquire(cap int, genPeriod time.Duration, genCount int, take int) *TokenBucketState {
	b.m.Lock()
	defer b.m.Unlock()

	cap = check_positive(cap)
	genCount = check_positive(genCount)
	take = check_positive(take)
	interval := check_positive(genPeriod.Nanoseconds() / int64(genCount))
	now := time.Now().UnixNano()

	if now > b.t+interval {
		add := (now - b.t) / interval
		b.c += add
		b.t += add * interval
	}
	if b.c > int64(cap) {
		b.c = int64(cap)
	}

	state := &TokenBucketState{
		cap:      cap,
		period:   genPeriod,
		count:    genCount,
		interval: interval,
		take:     take,
		now:      now,
		remain:   b.c,
		genTime:  b.t,
	}

	if b.c < int64(take) {
		return state
	}
	b.c -= int64(take)
	state.result = true
	state.remain = b.c
	return state
}

func check_positive[T int | int64](n T) T {
	if n < 1 {
		return 1
	}
	return n
}

type TokenBucketState struct {
	cap      int           // 令牌桶容量
	period   time.Duration // 令牌生产周期
	count    int           // 令牌周期内生产数量
	interval int64         // 每个令牌生产周期
	take     int
	now      int64
	result   bool  // 获取令牌结果
	remain   int64 // 桶内剩余令牌
	genTime  int64 // 最近生产令牌时间
}

func (s *TokenBucketState) OK() bool {
	return s.result
}

func (s *TokenBucketState) Remaining() int {
	return int(s.remain)
}

func (s *TokenBucketState) NextGenAfter() time.Duration {
	return time.Duration(s.genTime + s.interval - s.now)
}
